# this script should replicate replicate a lattice planner
# algorithm follow:
# 1. get reference point
# 2. transform ego location from Cartesian to frenet
# 3. longitude sample t->s (4?/5?)
# 4. lateral sample s->l
# 5. filter invalid longitude and lateral traj
# 6. calc cost, sort in priority_queue
# 7. merge traj
# 8. collison detect
# controller
# 9. Translation and rotation
# 10. update ego position, return to step 1
import heapq

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from math import *
from scipy import optimize


def calcDist(pose1,pose2):
    return sqrt((pose1[0]-pose2[0])**2+(pose1[1]-pose2[1])**2)

class RefLine:
    def __init__(self,x_data,y_data):
        self.f=np.poly1d(np.polyfit(x_data,y_data,len(x_data)-1))
    def getRefLine(self):
        return self.f

class GetRefPoint:
    def __init__(self,refline,xmin,xmax):
        self.refline_table={"x":[],"y":[],"s":[]}
        self.refline=refline
        for x in np.linspace(xmin,xmax,100):
            self.refline_table["x"].append(x)
            t_y=refline.getRefLine()(x)
            self.refline_table["y"].append(t_y)
            if len(self.refline_table["s"])==0:
                self.refline_table["s"].append(0)
            else:
                xk1=self.refline_table["x"][-1]
                yk1=self.refline_table["y"][-1]
                sk1=self.refline_table["s"][-1]
                t_s=calcDist([x,t_y],[xk1,yk1])+sk1
                self.refline_table["s"].append(t_s)

    def getCloestPoint(self,ego_pose):
        temp=999
        cloest_index=0
        for index in range(len(self.refline_table["x"])):
            t_x=self.refline_table["x"][index]
            t_y=self.refline_table["y"][index]
            t_dis=calcDist(ego_pose,[t_x,t_y])
            if t_dis<temp:
                temp=t_dis
                cloest_index=index
            elif t_y>=1.5*ego_pose[0]:
                break
        rs=self.refline_table["s"][cloest_index]
        rx=self.refline_table["x"][cloest_index]
        ry=self.refline_table["y"][cloest_index]
        rtheta=np.polyder(self.refline)(rx)
        rkappa=np.polyder(self.refline,2)(rx)
        rdkappa=np.polyder(self.refline_table,3)(rx)

        return [rs,rx,ry,rtheta,rkappa,rdkappa]
    def getRefTable(self):
        return self.refline_table

# store ego historical point, calc kappa and theta
class Ego_property:
    def __init__(self,length):
        self.ego_point_list=[]
        self.length=length
        self.Traj=np.poly1d
    def addPoint(self,pose):
        if (len(self.ego_point_list)>=self.length):
            self.ego_point_list.pop(0)
        self.ego_point_list.append(pose)
    def updateTraj(self):
        x=[]
        y=[]
        for item in self.ego_point_list:
            x.append(item[0])
            y.append(item[1])
        f=np.polyfit(x,y,3)
        self.Traj=np.poly1d(f)
    def getEgoProperty(self,pose,v,a):
        x=pose[0]
        y=pose[1]
        theta=np.polyder(self.Traj)(x)
        kappa=np.polyder(self.Traj,2)(x)
        return [x,y,v,a,theta,kappa]

class C2F:
    def __init__(self,refpoint,egopoint):
        rs=refpoint[0]
        rx=refpoint[1]
        ry=refpoint[2]
        rtheta=refpoint[3]
        rkappa=refpoint[4]
        rdkappa=refpoint[5]
        x=egopoint[0]
        y=egopoint[1]
        v=egopoint[2]
        a=egopoint[3]
        theta=egopoint[4]
        kappa=egopoint[5]
        s_condition = np.zeros(3)
        d_condition = np.zeros(3)

        dx = x - rx
        dy = y - ry

        cos_theta_r = cos(rtheta)
        sin_theta_r = sin(rtheta)

        cross_rd_nd = cos_theta_r * dy - sin_theta_r * dx
        d_condition[0] = copysign(sqrt(dx * dx + dy * dy), cross_rd_nd)

        delta_theta = theta - rtheta
        tan_delta_theta = tan(delta_theta)
        cos_delta_theta = cos(delta_theta)

        one_minus_kappa_r_d = 1 - rkappa * d_condition[0]
        d_condition[1] = one_minus_kappa_r_d * tan_delta_theta

        kappa_r_d_prime = rdkappa * d_condition[0] + rkappa * d_condition[1]

        d_condition[2] = (-kappa_r_d_prime * tan_delta_theta +
                          one_minus_kappa_r_d / cos_delta_theta / cos_delta_theta *
                          (kappa * one_minus_kappa_r_d / cos_delta_theta - rkappa))

        s_condition[0] = rs
        s_condition[1] = v * cos_delta_theta / one_minus_kappa_r_d

        delta_theta_prime = one_minus_kappa_r_d / cos_delta_theta * kappa - rkappa
        s_condition[2] = ((a * cos_delta_theta -
                           s_condition[1] * s_condition[1] *
                           (d_condition[1] * delta_theta_prime - kappa_r_d_prime)) /
                          one_minus_kappa_r_d)
        self.s_condition=s_condition
        self.d_condition=d_condition
    def get(self):
        return [self.s_condition,self.d_condition]


class F2C:
    def run(self,rs, rx, ry, rtheta, s_condition, d_condition):
        if fabs(rs - s_condition[0]) >= 1.0e-6:
            print("The reference point s and s_condition[0] don't match")

        cos_theta_r = cos(rtheta)
        sin_theta_r = sin(rtheta)

        x = rx - sin_theta_r * d_condition[0]
        y = ry + cos_theta_r * d_condition[0]

        return x, y


class QuinticSample:# is used for lateral and longitude sample
    def __init__(self,origion,end,nums,var_array):
        self.origion=origion
        self.end=end
        self.nums=nums
        self.var_array=var_array

    def run(self):
        p_range=self.end["p"]
        v_range=self.end["v"]
        a_range=self.end["a"]

        p_array=np.arange(p_range[0],p_range[1],self.nums)
        v_array=np.arange(v_range[0],v_range[1],self.nums)
        a_array=np.arange(a_range[0],a_range[1],self.nums)
        finalstate_array = []
        for i in range(self.nums):
            finalstate_array.append([p_array[i],v_array[i],a_array[i]])
        coeff_list=[]
        for t in self.var_array:
            for final_state in finalstate_array:
                O = self.origion + final_state
                M = np.array([
                    [1, 0, 0, 0, 0, 0],
                    [0, 1, 0, 0, 0, 0],
                    [0, 0, 2, 0, 0, 0],
                    [1, t, t ** 2, t ** 3, t ** 4, t ** 5],
                    [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3, 5 * t ** 4],
                    [0, 0, 2, 6 * t, 12 * t ** 2, 20 * t ** 3]
                ])
                O = np.array([O]).T
                t_co = np.dot(np.linalg.inv(M), O)
                t_co = np.flipud(t_co).flatten()
                t_fun = np.poly1d(t_co.astype(float))
                coeff_list.append({"var":t,"f":t_fun})
        return coeff_list

class QuarticSample:# just for long cruise sample
    def __init__(self,origion,end,nums,steptimes):
        self.origion=origion
        self.end=end
        self.nums=nums
        self.steptimes=steptimes
    def run(self):
        v_array = np.linspace(self.end["v"][0], self.end["v"][1], self.nums)
        a_array = np.linspace(self.end["a"][0], self.end["a"][1], self.nums)
        finalstate_array = []
        for i in range(self.nums):
            finalstate_array.append([v_array[i], a_array[i]])

        coeff_list = []
        for t in self.steptimes:
            for final_longitude_state in finalstate_array:
                O = [0] + self.origion + final_longitude_state
                M = [
                    [1, 0, 0, 0, 0],  # 初始位移
                    [0, 1, 0, 0, 0],  # 初始速度
                    [0, 0, 2, 0, 0],  # 初始加速度
                    [0, 1, 2 * t, 3 * t ** 2, 4 * t ** 3],
                    [0, 0, 2, 6 * t, 12 * t ** 2]
                ]
                O = np.array([O]).T
                t_co = np.dot(np.linalg.inv(M), O)
                t_co = np.flipud(t_co).flatten()
                t_fun = np.poly1d(t_co.astype(float))
                coeff_list.append({"var":t,"f":t_fun})
        return coeff_list

class Filter:
    def __init__(self,lateralTraj,longTraj):
        self.lateralTraj=lateralTraj
        self.longTraj=longTraj

    def run(self):
        for index in range(len(self.longTraj)-1,-1,-1):
            temp_f=self.longTraj[index]["f"]
            end_t=self.longTraj[index]["var"]

            vf=np.polyder(temp_f)
            vmin=vf(optimize.fminbound(vf,0,end_t))
            vmax=vf(optimize.fminbound(lambda x:-vf(x),0,end_t))
            af=np.polyder(temp_f,2)
            amin=af(optimize.fminbound(af,0,end_t))
            amax=af(optimize.fminbound(lambda x:-af(x),0,end_t))
            jerkf=np.polyder(temp_f,3)
            jerkmin=jerkf(optimize.fminbound(jerkf,0,end_t))
            jerkmax=jerkf(optimize.fminbound(lambda x: -jerkf(x),0,end_t))

            if (vmin<-1 or vmax>10
                or amin<-4.5 or amax>4
                or jerkmin<-4 or jerkmax>2):
                self.longTraj[index].pop(index)

class CostFun:
    def __init__(self,lateralTraj,longTraj):
        self.lateralTraj=lateralTraj
        self.longTraj=longTraj
        self.TrajSortByCost=[]
    def run(self):
        for laterT in self.lateralTraj:
            for longT in self.longTraj:
                total_dis=0
                s_list=np.arange(0,laterT["var"],0.5)
                for s in s_list:
                    total_dis += laterT["f"](s)
                cost = total_dis / len(s_list)
                tTraj=[cost,laterT,longT]
                self.TrajSortByCost.append(tTraj)
        self.TrajSortByCost.sort(key=lambda x:(x[0]))
    def getTop(self):
        tTraj=self.TrajSortByCost[0]
        self.TrajSortByCost.pop(0)
        return tTraj


class Merge:
    def __init__(self,Traj):
        self.Traj=Traj
        self.TrajPoint=[]
    def run(self,tTable,refline_fun):
        t=0.1 # predict time
        s = self.Traj[2]["f"](t)
        l = self.Traj[1]["f"](s)
        start = 0
        end = len(tTable) - 1
        mid = (start + end) / 2
        targetpoint = 0
        while start <= end:
            if (tTable[mid]["s"] > s):
                end = mid - 1
            elif tTable[mid]["s"] < s:
                start = mid + 1
            else:
                start = mid
        rx = tTable[start]["x"]
        ry = tTable[start]["y"]
        rs = tTable[start]["s"]
        rtheta = np.polyder(refline_fun)(rx)
        return [rs,rx,ry,rtheta,s,l]




class CollisonDetect:
    pass

class MoveIt:
    pass

if __name__ == '__main__':
    # plot env
    x=np.array([10,12,15,16])
    y=np.array([0,45,90,100])

    refline=RefLine(x,y)
    x_data=np.linspace(9,16,100)
    y_data=refline.getRefLine()(x_data)
    print(refline.getRefLine())

    plt.plot(x_data,y_data)
    # plot obj
    x=[11,15]
    y=[45,80]
    plt.scatter(x,y,c='r')
    # plot ego
    ego_pose=[9,1]
    plt.scatter(ego_pose[0],ego_pose[1],c='g')

    # start lattice
    RefPoint=GetRefPoint(refline,10,20)
    cloest_point = RefPoint.getCloestPoint(ego_pose)
    # end
    plt.xlim(0,25)
    plt.ylim(0,120)
    plt.show()

